package com.xu.cloud.common.trace;

import org.springframework.core.NamedThreadLocal;
import org.springframework.lang.NonNull;
import org.springframework.lang.Nullable;

import java.util.Objects;

/**
 * 本地线程缓存TraceId上下文
 *
 * @author xuguan
 * @since 2025/1/8
 */
public class TraceIdContextHolder {
    private static final ThreadLocal<TraceIdContext> traceIdContextHolder = new NamedThreadLocal<>("TraceIdContext");

    @Nullable
    public static TraceIdContext getTraceIdContext() {
        return traceIdContextHolder.get();
    }

    @NonNull
    public static TraceIdContext getOrSetTraceIdContext() {
        TraceIdContext traceIdContext = traceIdContextHolder.get();
        if (Objects.isNull(traceIdContext)) {
            traceIdContext = new UUIDTraceIdContext();
        }
        return traceIdContext;
    }

    public static void setTraceIdContext(TraceIdContext traceIdContext) {
        traceIdContextHolder.set(traceIdContext);
    }

    public static void removeTraceIdContext() {
        traceIdContextHolder.remove();
    }
}
